import random
from math import sqrt
from tkinter import *
from typing import List, Tuple, Optional, Union

import gym
import numpy as np


class GridWorldEnv(gym.core.Env):
    GRID_W = 80
    W, H = 9, 7
    MOVE_UP, MOVE_DOWN, MOVE_LEFT, MOVE_RIGHT = range(4)
    TYPE_GRID, TYPE_WALL, TYPE_TRAP, TYPE_GOAL, TYPE_PLAYER = range(5)
    MAX_STEP = 500  # 最多走500步

    walls = [(1, 1), (2, 2), (2, 1), (6, 1), (7, 2), (7, 1), (3, 5), (4, 5), (5, 5)]
    specials = [(1, 2, "red", -600), (6, 2, "red", -600), (4, 3, "red", -600), (4, 0, "green", 100)]
    action_space = gym.spaces.Discrete(4)
    observation_space = gym.spaces.Box(
        low=0, high=255, shape=(W, H), dtype=np.int
    )

    def __init__(self):
        self.master = Tk()
        self.board = Canvas(self.master, width=self.W * self.GRID_W, height=self.H * self.GRID_W)
        self.grid_map, self.player = self.init()
        self.start_point = self.player
        self.player_rect = self.board.create_rectangle(
            *[self.player[i] * self.GRID_W + self.GRID_W * j / 10 for j in (2, 8) for i in range(2)],
            fill="orange", width=1, tag="player"
        )
        self.board.grid(row=0, column=0)
        self.total_step = 0
        self.goal = (4, 0)

    def init(self) -> (np.ndarray, List[int]):
        grid_map = []
        for row in range(self.H):
            for col in range(self.W):
                self.board.create_rectangle(
                    *[(i + j) * self.GRID_W for j in range(2) for i in (col, row)],
                    fill="white", width=1
                )
            grid_map.append([self.TYPE_GRID] * self.W)
        for (col, row, color, rew) in self.specials:
            self.board.create_rectangle(
                *[(i + j) * self.GRID_W for j in range(2) for i in (col, row)],
                fill=color, width=1
            )
            grid_map[row][col] = self.TYPE_TRAP if rew < 0 else self.TYPE_GOAL
        for (col, row) in self.walls:
            self.board.create_rectangle(
                *[(i + j) * self.GRID_W for j in range(2) for i in (col, row)],
                fill="black", width=1
            )
            grid_map[row][col] = self.TYPE_WALL

        player = [int(self.W / 2), self.H - 1]
        grid_map[player[1]][player[0]] = self.TYPE_PLAYER
        return np.array(grid_map), player

    def step(self, action: int):
        # print(self.get_action_meanings()[action])
        reward, done = self.move(
            dx=1 if action == self.MOVE_RIGHT else (-1 if action == self.MOVE_LEFT else 0),
            dy=1 if action == self.MOVE_DOWN else (-1 if action == self.MOVE_UP else 0)
        )
        self.total_step += 1
        if self.total_step >= self.MAX_STEP:
            done = True
        return self.grid_map, reward, done, {}

    def move(self, dx, dy) -> (int, bool):
        self.grid_map[self.player[1]][self.player[0]] = self.TYPE_GRID
        player = [self.player[0] + dx, self.player[1] + dy]
        if (0 <= player[0] < self.W) and \
                (0 <= player[1] < self.H) and \
                ((player[0], player[1]) not in self.walls):
            self.board.coords(
                self.player_rect,
                *[player[i] * self.GRID_W + self.GRID_W * j / 10 for j in (2, 8) for i in range(2)]
            )
            self.player = player
        self.grid_map[self.player[1]][self.player[0]] = self.TYPE_PLAYER
        for (i, j, c, w) in self.specials:
            if self.player[0] == i and self.player[1] == j:
                return w, True
        return -1, False

    def reset(self):
        self.close()
        self.__init__()
        return self.grid_map

    def close(self):
        self.master.destroy()

    def seed(self, seed=None):
        if seed is not None:
            random.seed(seed)
            return [seed]
        else:
            return []

    def render(self, mode='human'):
        try:
            self.master.update()
        except TclError:
            pass

    @staticmethod
    def get_action_meanings():
        return ["UP", "DOWN", "LEFT", "RIGHT"]

    @staticmethod
    def manhattan(p1: Union[Tuple[int, int], List[int]], p2: Union[Tuple[int, int], List[int]]):
        # 两点的曼哈顿距离
        return abs(p1[0] - p2[0]) + abs(p1[1] - p2[1])

    @staticmethod
    def distance(p1: Union[Tuple[int, int], List[int]], p2: Union[Tuple[int, int], List[int]]):
        # 两点的直线距离
        return sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)